//
// Created by xiang on 2021/7/19.
//

#include "ch4/gins_pre_integ.h"
#include "ch4/g2o_types.h"
#include "common/g2o_types.h"

#include <glog/logging.h>

#include <g2o/core/block_solver.h>
#include <g2o/core/optimization_algorithm_levenberg.h>
#include <g2o/core/robust_kernel.h>
#include <g2o/solvers/eigen/linear_solver_eigen.h>

namespace sad {

void GinsPreInteg::AddImu(const IMU& imu) {
    if (first_gnss_received_ && first_imu_received_) {
        pre_integ_->Integrate(imu, imu.timestamp_ - last_imu_.timestamp_);
    }

    //LOG(INFO) << "imu.timestamp_: " << imu.timestamp_ << ", " << last_imu_.timestamp_;

    first_imu_received_ = true;
    last_imu_ = imu;
    current_time_ = imu.timestamp_;
}

void GinsPreInteg::SetOptions(sad::GinsPreInteg::Options options) {
    double bg_rw2 = 1.0 / (options_.bias_gyro_var_ * options_.bias_gyro_var_);
    options_.bg_rw_info_.diagonal() << bg_rw2, bg_rw2, bg_rw2;
    double ba_rw2 = 1.0 / (options_.bias_acce_var_ * options_.bias_acce_var_);
    options_.ba_rw_info_.diagonal() << ba_rw2, ba_rw2, ba_rw2;

    double gp2 = options_.gnss_pos_noise_ * options_.gnss_pos_noise_;
    double gh2 = options_.gnss_height_noise_ * options_.gnss_height_noise_;
    double ga2 = options_.gnss_ang_noise_ * options_.gnss_ang_noise_;

    options_.gnss_info_.diagonal() << 1.0 / ga2, 1.0 / ga2, 1.0 / ga2, 1.0 / gp2, 1.0 / gp2, 1.0 / gh2;
    pre_integ_ = std::make_shared<IMUPreintegration>(options_.preinteg_options_);

    double o2 = 1.0 / (options_.odom_var_ * options_.odom_var_);
    options_.odom_info_.diagonal() << o2, o2, o2;

    prior_info_.block<6, 6>(9, 9) = Mat6d ::Identity() * 1e6;

    if (this_frame_) {
        this_frame_->bg_ = options_.preinteg_options_.init_bg_;
        this_frame_->ba_ = options_.preinteg_options_.init_ba_;
    }
}

void GinsPreInteg::AddGnss(const GNSS& gnss) {
    this_frame_ = std::make_shared<NavStated>(current_time_);
    this_gnss_ = gnss;

    if (!first_gnss_received_) {
        if (!gnss.heading_valid_) {
            // 要求首个GNSS必须有航向
            return;
        }

        // 首个gnss信号，将初始pose设置为该gnss信号
        this_frame_->timestamp_ = gnss.unix_time_;
        this_frame_->p_ = gnss.utm_pose_.translation();
        this_frame_->R_ = gnss.utm_pose_.so3();
        this_frame_->v_.setZero();
        this_frame_->bg_ = options_.preinteg_options_.init_bg_;
        this_frame_->ba_ = options_.preinteg_options_.init_ba_;

        pre_integ_ = std::make_shared<IMUPreintegration>(options_.preinteg_options_);

        keyframe_ = std::make_shared<KeyFrame>(this_frame_, pre_integ_, this_gnss_);
        keyframe_queues_.push_back(keyframe_);
        
        last_frame_ = this_frame_;
        last_gnss_ = this_gnss_;
        first_gnss_received_ = true;
        current_time_ = gnss.unix_time_;
        return;
    }

    // 积分到GNSS时刻
    pre_integ_->Integrate(last_imu_, gnss.unix_time_ - current_time_);

    //LOG(INFO) << "gnss.unix_time_: " << gnss.unix_time_ << ", " << current_time_;

    current_time_ = gnss.unix_time_;
    *this_frame_ = pre_integ_->Predict(*last_frame_, options_.gravity_);
    
    keyframe_ = std::make_shared<KeyFrame>(this_frame_, pre_integ_, this_gnss_);
    keyframe_queues_.push_back(keyframe_);
    
    Optimize();
    
    //GlobalOptimize();
    
    last_frame_ = this_frame_;
    last_gnss_ = this_gnss_;
}

void GinsPreInteg::AddOdom(const sad::Odom& odom) {
    last_odom_ = odom;
    last_odom_set_ = true;
}

void GinsPreInteg::Optimize() {
    if (pre_integ_->dt_ < 1e-3) {
        // 未得到积分
        return;
    }

    using BlockSolverType = g2o::BlockSolverX;
    using LinearSolverType = g2o::LinearSolverEigen<BlockSolverType::PoseMatrixType>;

    auto* solver = new g2o::OptimizationAlgorithmLevenberg(
    g2o::make_unique<BlockSolverType>(g2o::make_unique<LinearSolverType>()));
    g2o::SparseOptimizer optimizer;
    optimizer.setAlgorithm(solver);

    // 上时刻顶点， pose, v, bg, ba
    auto v0_pose = new VertexPose();
    v0_pose->setId(0);
    v0_pose->setEstimate(last_frame_->GetSE3());
    optimizer.addVertex(v0_pose);

    auto v0_vel = new VertexVelocity();
    v0_vel->setId(1);
    v0_vel->setEstimate(last_frame_->v_);
    optimizer.addVertex(v0_vel);

    auto v0_bg = new VertexGyroBias();
    v0_bg->setId(2);
    v0_bg->setEstimate(last_frame_->bg_);
    optimizer.addVertex(v0_bg);

    auto v0_ba = new VertexAccBias();
    v0_ba->setId(3);
    v0_ba->setEstimate(last_frame_->ba_);
    optimizer.addVertex(v0_ba);

    // 本时刻顶点，pose, v, bg, ba
    auto v1_pose = new VertexPose();
    v1_pose->setId(4);
    v1_pose->setEstimate(this_frame_->GetSE3());
    optimizer.addVertex(v1_pose);

    auto v1_vel = new VertexVelocity();
    v1_vel->setId(5);
    v1_vel->setEstimate(this_frame_->v_);
    optimizer.addVertex(v1_vel);

    auto v1_bg = new VertexGyroBias();
    v1_bg->setId(6);
    v1_bg->setEstimate(this_frame_->bg_);
    optimizer.addVertex(v1_bg);

    auto v1_ba = new VertexAccBias();
    v1_ba->setId(7);
    v1_ba->setEstimate(this_frame_->ba_);
    optimizer.addVertex(v1_ba);

    // 预积分边
    auto edge_inertial = new EdgeInertial(pre_integ_, options_.gravity_);
    edge_inertial->setVertex(0, v0_pose);
    edge_inertial->setVertex(1, v0_vel);
    edge_inertial->setVertex(2, v0_bg);
    edge_inertial->setVertex(3, v0_ba);
    edge_inertial->setVertex(4, v1_pose);
    edge_inertial->setVertex(5, v1_vel);
    auto* rk = new g2o::RobustKernelHuber();
    rk->setDelta(200.0);
    edge_inertial->setRobustKernel(rk);
    optimizer.addEdge(edge_inertial);

    // 两个零偏随机游走
    auto* edge_gyro_rw = new EdgeGyroRW();
    edge_gyro_rw->setVertex(0, v0_bg);
    edge_gyro_rw->setVertex(1, v1_bg);
    edge_gyro_rw->setInformation(options_.bg_rw_info_);
    optimizer.addEdge(edge_gyro_rw);

    auto* edge_acc_rw = new EdgeAccRW();
    edge_acc_rw->setVertex(0, v0_ba);
    edge_acc_rw->setVertex(1, v1_ba);
    edge_acc_rw->setInformation(options_.ba_rw_info_);
    optimizer.addEdge(edge_acc_rw);

    // 上时刻先验
    auto* edge_prior = new EdgePriorPoseNavState(*last_frame_, prior_info_);
    edge_prior->setVertex(0, v0_pose);
    edge_prior->setVertex(1, v0_vel);
    edge_prior->setVertex(2, v0_bg);
    edge_prior->setVertex(3, v0_ba);
    optimizer.addEdge(edge_prior);

    // GNSS边
    auto edge_gnss0 = new EdgeGNSS(v0_pose, last_gnss_.utm_pose_);
    edge_gnss0->setInformation(options_.gnss_info_);
    optimizer.addEdge(edge_gnss0);

    auto edge_gnss1 = new EdgeGNSS(v1_pose, this_gnss_.utm_pose_);
    edge_gnss1->setInformation(options_.gnss_info_);
    optimizer.addEdge(edge_gnss1);

    // Odom边
    EdgeEncoder3D* edge_odom = nullptr;
    Vec3d vel_world = Vec3d::Zero();
    Vec3d vel_odom = Vec3d::Zero();
    if (last_odom_set_) {
        // velocity obs
        double velo_l =
            options_.wheel_radius_ * last_odom_.left_pulse_ / options_.circle_pulse_ * 2 * M_PI / options_.odom_span_;
        double velo_r =
            options_.wheel_radius_ * last_odom_.right_pulse_ / options_.circle_pulse_ * 2 * M_PI / options_.odom_span_;
        double average_vel = 0.5 * (velo_l + velo_r);
        vel_odom = Vec3d(average_vel, 0.0, 0.0);
        vel_world = this_frame_->R_ * vel_odom;

        edge_odom = new EdgeEncoder3D(v1_vel, vel_world);
        edge_odom->setInformation(options_.odom_info_);
        optimizer.addEdge(edge_odom);

        // 重置odom数据到达标志位，等待最新的odom数据
        last_odom_set_ = false;
    }

    optimizer.setVerbose(options_.verbose_);
    optimizer.initializeOptimization();
    optimizer.optimize(20);

    if (options_.verbose_) {
        // 获取结果，统计各类误差
        LOG(INFO) << "chi2/error: ";
        LOG(INFO) << "preintegration: " << edge_inertial->chi2() << "/" << edge_inertial->error().transpose();
        // LOG(INFO) << "gnss0: " << edge_gnss0->chi2() << ", " << edge_gnss0->error().transpose();
        LOG(INFO) << "gnss1: " << edge_gnss1->chi2() << ", " << edge_gnss1->error().transpose();
        LOG(INFO) << "bias: " << edge_gyro_rw->chi2() << "/" << edge_acc_rw->error().transpose();
        LOG(INFO) << "prior: " << edge_prior->chi2() << "/" << edge_prior->error().transpose();
        if (edge_odom) {
            LOG(INFO) << "body vel: " << (v1_pose->estimate().so3().inverse() * v1_vel->estimate()).transpose();
            LOG(INFO) << "meas: " << vel_odom.transpose();
            LOG(INFO) << "odom: " << edge_odom->chi2() << "/" << edge_odom->error().transpose();
        }
    }
    // useless set last frame, because will be update later
    last_frame_->R_ = v0_pose->estimate().so3();
    last_frame_->p_ = v0_pose->estimate().translation();
    last_frame_->v_ = v0_vel->estimate();
    last_frame_->bg_ = v0_bg->estimate();
    last_frame_->ba_ = v0_ba->estimate();

    this_frame_->R_ = v1_pose->estimate().so3();
    this_frame_->p_ = v1_pose->estimate().translation();
    this_frame_->v_ = v1_vel->estimate();
    this_frame_->bg_ = v1_bg->estimate();
    this_frame_->ba_ = v1_ba->estimate();

    // 重置integ
    options_.preinteg_options_.init_bg_ = this_frame_->bg_;
    options_.preinteg_options_.init_ba_ = this_frame_->ba_;
    pre_integ_ = std::make_shared<IMUPreintegration>(options_.preinteg_options_);

    // 计算当前时刻先验
    // 构建hessian
    // 15x2，顺序：v0_pose, v0_vel, v0_bg, v0_ba, v1_pose, v1_vel, v1_bg, v1_ba
    //            0       6        9     12     15        21      24     27
    Eigen::Matrix<double, 30, 30> H;
    H.setZero();

    H.block<24, 24>(0, 0) += edge_inertial->GetHessian();

    Eigen::Matrix<double, 6, 6> Hgr = edge_gyro_rw->GetHessian();
    H.block<3, 3>(9, 9) += Hgr.block<3, 3>(0, 0);
    H.block<3, 3>(9, 24) += Hgr.block<3, 3>(0, 3);
    H.block<3, 3>(24, 9) += Hgr.block<3, 3>(3, 0);
    H.block<3, 3>(24, 24) += Hgr.block<3, 3>(3, 3);

    Eigen::Matrix<double, 6, 6> Har = edge_acc_rw->GetHessian();
    H.block<3, 3>(12, 12) += Har.block<3, 3>(0, 0);
    H.block<3, 3>(12, 27) += Har.block<3, 3>(0, 3);
    H.block<3, 3>(27, 12) += Har.block<3, 3>(3, 0);
    H.block<3, 3>(27, 27) += Har.block<3, 3>(3, 3);

    H.block<15, 15>(0, 0) += edge_prior->GetHessian();
    H.block<6, 6>(0, 0) += edge_gnss0->GetHessian();
    H.block<6, 6>(15, 15) += edge_gnss1->GetHessian();
    if(edge_odom != nullptr) {
       H.block<3, 3>(6, 6) += edge_odom->GetHessian();
    }

    H = math::Marginalize(H, 0, 14);
    prior_info_ = H.block<15, 15>(15, 15);
    
    //LOG(INFO) << "H.block<15, 15>(0, 0): " << H.block<15, 15>(0, 0);
    //LOG(INFO) << "H.block<15, 15>(15, 15): " << H.block<15, 15>(15, 15);
}

void GinsPreInteg::GlobalOptimize() {
     if(keyframe_queues_.size() < 2) {
        return;
     }
     
     //if(keyframe_queues_.size() > 10) {
     //   keyframe_queues_.pop_front();
     //}

     using BlockSolverType = g2o::BlockSolverX;
     using LinearSolverType = g2o::LinearSolverEigen<BlockSolverType::PoseMatrixType>;

     auto *solver = new g2o::OptimizationAlgorithmLevenberg(
     g2o::make_unique<BlockSolverType>(g2o::make_unique<LinearSolverType>()));
     g2o::SparseOptimizer optimizer;
     optimizer.setAlgorithm(solver);
   
     int max_keframe_idx = -1;
     std::vector<VertexPose*> vertex_pose;
     std::vector<VertexVelocity*> vertex_vel;
     std::vector<VertexGyroBias*> vertex_bg;
     std::vector<VertexAccBias*> vertex_ba;
     int keyframe_idx = 0;
     std::deque<std::shared_ptr<KeyFrame>>::iterator iter; 
     for(iter = keyframe_queues_.begin(); iter != keyframe_queues_.end(); ++iter) {
        int current_id = keyframe_idx;
        keyframe_idx++;
        //int current_id = (*iter)->id_;

        //if(current_id < 5)
        //    std::cout << " current_id : " << current_id << " , " << (*iter)->id_ << std::endl;

        if((current_id) > max_keframe_idx) {
            max_keframe_idx = (current_id);
        }

        VertexPose* pose = new VertexPose();
        pose->setId(4*current_id);
        pose->setEstimate((*iter)->this_frame()->GetSE3());
        if(current_id == 0) {
            pose->setFixed(true);
        }
        optimizer.addVertex(pose);

        VertexVelocity* vel = new VertexVelocity();
        vel->setId(4*current_id + 1);
        vel->setEstimate((*iter)->this_frame()->v_);
        if(current_id == 0) {
            vel->setFixed(true);
        }
        optimizer.addVertex(vel);

        VertexGyroBias* bg = new VertexGyroBias();
        bg->setId(4*current_id + 2);
        bg->setEstimate((*iter)->this_frame()->bg_);
        if(current_id == 0) {
            bg->setFixed(true);
        }
        optimizer.addVertex(bg);

        VertexAccBias* ba = new VertexAccBias();
        ba->setId(4*current_id + 3);
        ba->setEstimate((*iter)->this_frame()->ba_);
        if(current_id == 0) {
            ba->setFixed(true);
        }
        optimizer.addVertex(ba);

        vertex_pose.push_back(pose);
        vertex_vel.push_back(vel);
        vertex_bg.push_back(bg);
        vertex_ba.push_back(ba);
     }

     int idx = 0;
     std::deque<std::shared_ptr<KeyFrame>>::iterator iter_edge; 
     for(iter_edge = keyframe_queues_.begin(); iter_edge != keyframe_queues_.end(); ++iter_edge) {
        if(iter_edge == keyframe_queues_.begin()) {
            continue;
        }

        std::shared_ptr<KeyFrame> keyframe = (*iter_edge);
        // imu factor
        EdgeInertial* edge_inertial = new EdgeInertial(keyframe->GetPreintegration(), options_.gravity_);
        edge_inertial->setVertex(0, vertex_pose[idx-1]);
        edge_inertial->setVertex(1, vertex_vel[idx-1]);
        edge_inertial->setVertex(2, vertex_bg[idx-1]);
        edge_inertial->setVertex(3, vertex_ba[idx-1]);
        edge_inertial->setVertex(4, vertex_pose[idx]);
        edge_inertial->setVertex(5, vertex_vel[idx]);
        auto *rk = new g2o::RobustKernelHuber();
        rk->setDelta(200.0);
        edge_inertial->setRobustKernel(rk);
        optimizer.addEdge(edge_inertial);        

        idx++;
     }

     for(int vi = 1; vi < vertex_bg.size(); vi++) {
        auto* edge_gyro_rw = new EdgeGyroRW();
        edge_gyro_rw->setVertex(0, vertex_bg[vi-1]);
        edge_gyro_rw->setVertex(1, vertex_bg[vi]);
        edge_gyro_rw->setInformation(options_.bg_rw_info_);
        optimizer.addEdge(edge_gyro_rw);        
     }

     for(int vi = 1; vi < vertex_ba.size(); vi++) {
        auto* edge_acc_rw = new EdgeAccRW();
        edge_acc_rw->setVertex(0, vertex_ba[vi-1]);
        edge_acc_rw->setVertex(1, vertex_ba[vi]);
        edge_acc_rw->setInformation(options_.ba_rw_info_);
        optimizer.addEdge(edge_acc_rw);        
     }

 
     for(int vp = 1; vp < keyframe_queues_.size(); vp++) {
        auto* edge_prior = new EdgePriorPoseNavState(*(keyframe_queues_[vp-1]->this_frame()), prior_info_);
        edge_prior->setVertex(0, vertex_pose[vp-1]);
        edge_prior->setVertex(1, vertex_vel[vp-1]);
        edge_prior->setVertex(2, vertex_bg[vp-1]);
        edge_prior->setVertex(3, vertex_ba[vp-1]);
        optimizer.addEdge(edge_prior);
     }


     for(int vp = 1; vp < keyframe_queues_.size(); vp++) {
        auto edge_gnss0 = new EdgeGNSS(vertex_pose[vp-1], keyframe_queues_[vp-1]->this_gnss().utm_pose_);
        edge_gnss0->setInformation(options_.gnss_info_);
        optimizer.addEdge(edge_gnss0);

        auto edge_gnss1 = new EdgeGNSS(vertex_pose[vp], keyframe_queues_[vp]->this_gnss().utm_pose_);
        edge_gnss1->setInformation(options_.gnss_info_);
        optimizer.addEdge(edge_gnss1);
     }
    

     for(int vp = 0; vp < keyframe_queues_.size(); vp++) {
        std::shared_ptr<KeyFrame> keyframe = keyframe_queues_[vp];
        
        EdgeEncoder3D* edge_odom = nullptr;
        Vec3d vel_world = Vec3d::Zero();
        Vec3d vel_odom = Vec3d::Zero();
        if (last_odom_set_) {
            // velocity obs
            double velo_l =
                options_.wheel_radius_ * last_odom_.left_pulse_ / options_.circle_pulse_ * 2 * M_PI / options_.odom_span_;
            double velo_r =
                options_.wheel_radius_ * last_odom_.right_pulse_ / options_.circle_pulse_ * 2 * M_PI / options_.odom_span_;
            double average_vel = 0.5 * (velo_l + velo_r);
            vel_odom = Vec3d(average_vel, 0.0, 0.0);

            vel_world = keyframe->this_frame()->R_ * vel_odom;
            VertexVelocity* vel = vertex_vel[vp];
            edge_odom = new EdgeEncoder3D(vel, vel_world);
            edge_odom->setInformation(options_.odom_info_);
            optimizer.addEdge(edge_odom);

            // 重置odom数据到达标志位，等待最新的odom数据
            last_odom_set_ = false;
        }
     }
     // start optimize
     optimizer.setVerbose(false);
     optimizer.initializeOptimization();
     optimizer.optimize(10);

     for(int vp = 0; vp < keyframe_queues_.size(); vp++) {
        std::shared_ptr<KeyFrame> keyframe = keyframe_queues_[vp];
        
        //int current_id = keyframe->id_;
        //VertexPose* pose = static_cast<VertexPose*>(optimizer.vertex(4*current_id));
        //VertexVelocity* vel = optimizer.vertex(4*current_id+1);
        //VertexGyroBias* bg = optimizer.vertex(4*current_id+2);
        //VertexAccBias* ba = optimizer.vertex(4*current_id+3);

        VertexPose* pose = vertex_pose[vp];
        VertexVelocity* vel = vertex_vel[vp];
        VertexGyroBias* bg = vertex_bg[vp];
        VertexAccBias* ba = vertex_ba[vp];

        this_frame_->R_ = pose->estimate().so3();
        this_frame_->p_ = pose->estimate().translation();
        this_frame_->v_ = vel->estimate();
        this_frame_->bg_ = bg->estimate();
        this_frame_->ba_ = ba->estimate();

        keyframe->SetImuBias(this_frame_->bg_, this_frame_->ba_);
        keyframe->SetImuState(this_frame_);
     }

     //this_frame_->R_ = vertex_pose.back()->estimate().so3();
     //this_frame_->p_ = vertex_pose.back()->estimate().translation();
     //this_frame_->v_ = vertex_vel.back()->estimate();
     //this_frame_->bg_ = vertex_bg.back()->estimate();
     //this_frame_->ba_ = vertex_ba.back()->estimate();

     options_.preinteg_options_.init_bg_ = this_frame_->bg_;
     options_.preinteg_options_.init_ba_ = this_frame_->ba_;
     pre_integ_ = std::make_shared<IMUPreintegration>(options_.preinteg_options_);
}

NavStated GinsPreInteg::GetState() const {
    if (this_frame_ == nullptr) {
        return {};
    }

    if (pre_integ_ == nullptr) {
        return *this_frame_;
    }

    return pre_integ_->Predict(*this_frame_, options_.gravity_);
}

}  // namespace sad